package com.qupeng.demo.kafka.kafkaspringbootproducer.producer;

import org.apache.kafka.clients.producer.ProducerRecord;
import org.apache.kafka.common.header.internals.RecordHeaders;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.kafka.core.KafkaTemplate;
import org.springframework.kafka.support.SendResult;
import org.springframework.stereotype.Service;

import java.util.UUID;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import java.util.function.BiConsumer;

@Service
public class KafkaService<K, V> {

    @Autowired
    private KafkaTemplate<K, V> kafkaTemplate;

    public SendResult<K, V> sendToKafkaSync(final String topic, final K key, final V value) throws ExecutionException, InterruptedException {
        RecordHeaders recordHeaders = new RecordHeaders();
        String correlationId = UUID.randomUUID().toString();
        recordHeaders.add("correlation-id", correlationId.getBytes());
        final ProducerRecord<K, V> record = new ProducerRecord<>(topic, null, System.currentTimeMillis(), key, value, recordHeaders);
        CompletableFuture<SendResult<K, V>> future = kafkaTemplate.send(record);
        return future.get();
    }

    public void sendToKafkaAsync(String topic, final K key, final V value, BiConsumer<? super SendResult<K, V>, ? super Throwable> action) {
        RecordHeaders recordHeaders = new RecordHeaders();
        String correlationId = UUID.randomUUID().toString();
        recordHeaders.add("correlation-id", correlationId.getBytes());
        final ProducerRecord<K, V> record = new ProducerRecord<>(topic, null, System.currentTimeMillis(), key, value, recordHeaders);
        CompletableFuture<SendResult<K, V>> future = kafkaTemplate.send(record);
        future.whenCompleteAsync(action);
    }
}
